{-# LANGUAGE ScopedTypeVariables #-}
module Hometasks where


import Graph as G

import Data.Maybe

import System.Random.MWC

import Data.Vector.Generic as VG(drop, foldl, (!), slice)
import Data.Vector as V(Vector)

import Control.Monad.ST
import Control.Monad
import Data.List as L
import Data.Map as Map(Map,singleton,empty,insertWith,lookup)
import Data.Array.IArray as Arr
import Data.Array.Unboxed

task1 :: SimpleGraph -> [Int]
task1 g = getLongest (0,[]) $ -- и выбираем длинейший
          map (toLength . (\v -> findPath [v] g)) $ -- ищем для каждой оптимальный путь
          take (size g) $ -- строим список вершин в графе 
          iterate (\i -> i + 1) 1
    where getLongest (_,p) [] = p
          getLongest a@(l,p) (b@(l',p'):ps) = getLongest (if l' > l then b else a) ps

          toLength l = (length l,l) -- сопаставляем путю пару (длина пути, путь)

          findPath p@(i:_) g = 
              let adjacent = 
                     filter (\j -> not $ fromJust $ getWeight g i j) $ -- выбираем те, которые не были пройдены
                            getAdjacents g i -- получаем все смежные вершины
              in 
                 if null adjacent  -- если больше идти некуда
                    then p -- возвращаем текущий путь
                 else getLongest (length p, p) $ -- иначе ищем длинейший путь 
                         map (toLength . (\i' -> findPath (i':p) (updateBidirectEdge g i i' True))) adjacent -- среди путей продолженых в следующую вершину


-- Генерирует случайный граф(если @arg2 то ненаправленный) с @arg1 вершинами.
generateRandomGraph :: Variate a => Int -> (a -> a) -> Bool -> IO (Graph a)
generateRandomGraph n normalizer directed =
    if n < 0 then fail "Bad number of vertexies" else do
     es <- withSystemRandom $
                 \(gen::GenST s) -> uniformVector gen (n*n + 1) :: Variate a => ST s (Vector (Int, Int, a)) -- указываем, какой вектор хотим
     let eCount = flip mod (n*n) $ (\(a,_,_) -> a) $ es VG.! 0 -- число ребер берём из первого элемента
     print $ "Building graph with " ++ show eCount ++ " edges"
     let graph = VG.foldl 
                   (\g (i,j,val) -> 
                         (if directed then updateEdge else updateBidirectEdge)  -- если направленный - добавляем одно ребро, иначе оба
                            g (mod i n + 1) (mod j n + 1) (normalizer val)) -- от вершины i до вершины j с весом val
                     (isolatedVertixies n) $ -- на новом графе с n вершинами
                     VG.slice 1 eCount es     -- а номера вершин берём начиная со второго элемента до eCount
     
     return (Prelude.foldl (\g i -> delEdges g i i) graph [1..n])


minimalTree :: Graph Int -> [(Int,Int)]
minimalTree g = buildTree [] [] edges 
    where edges = sortBy (\(_,_,w1) (_,_,w2) -> compare w1 w2) $
                  deleteDuplicates $ 
                  sort $ 
                  concat $
                  map (\i -> map (\j -> (min i j, max i j,fromJust $ getWeight g i j)) $ 
                       getAdjacents g i) $ 
                  take (size g) [1,2..]
          deleteDuplicates [a] = [a]
          deleteDuplicates [] = []
          deleteDuplicates (x:y:es) 
              | x == y = deleteDuplicates (x:es)
              | otherwise = x : deleteDuplicates (y:es)
          buildTree _ edges [] = edges
          buildTree vert edges ((i,j,w):es)
              | not (elem i vert) ||
                not (elem j vert) = buildTree (i:j:vert) ((i,j):edges) es
              | otherwise = buildTree vert edges es 


travellingSalesman :: Graph Int -> Maybe (Int, [Int])
travellingSalesman g = do
        guard (all (not . null) $ map (getAdjacents g) $ [1..size g])
        let table = (array (1, size g) $ 
                [(i,Map.empty)| i <- [1..size g]]) // 
                    [(1, Map.singleton (array (1,size g) [(i,False)| i <- [1..size g]]) 0)]
            filledTable = 
                dynamic table (1, endPosition)
        cost <- Map.lookup endPosition $ filledTable Arr.! 0
        path <- rebuildPath [] filledTable 0 endPosition
        return (cost, path)
    where 
          endPosition = array (1,size g) [(i,i /= 1)| i <- [1..size g]]
          dynamic :: Array Int (Map (UArray Int Bool) Int) 
                -> (Int, UArray Int Bool) -> Array Int (Map (UArray Int Bool) Int)
          dynamic table (i, included) =
                L.foldl (\t j -> updateTable t i j included) table $
                    filter ((Arr.!) included) $ getAdjacents g i 
          updateTable :: Array Int (Map (UArray Int Bool) Int) 
                -> Int -> Int -> UArray Int Bool -> Array Int (Map (UArray Int Bool) Int)
          updateTable t i j included = 
              let currentMap = t Arr.! i 
                  included' = included // [(j,False)]
                  newMap = 
                    case Map.lookup included' $ dynamic t (j, included') Arr.! j of
                         Just v -> insertWith 
                                     min included 
                                     ((fromJust $ getWeight g i j) + v) 
                                     currentMap
                         Nothing -> currentMap
              in
                t // [(i, newMap)]
          rebuildPath :: [Int] -> Array Int (Map (UArray Int Bool) Int) -> Int -> UArray Int Bool -> Maybe [Int]
          rebuildPath p t i included 
                | all not $ elems included = Just p
                | otherwise = 
                    let currentValue = Map.lookup included $ t Arr.! i in 
                    case filter 
                          (\j -> areAdjacent g i j && included Arr.! j && 
                                  currentValue == maybeSum  
                                    (Map.lookup (included // [(j,False)]) $ t Arr.! j)
                                    (getWeight g i j) ) 
                          [j|j <- [1..size g]]
                         of
                     [j] -> rebuildPath (j:p) t j (included // [(j,False)])
                     [] -> Nothing
          maybeSum Nothing _ = Nothing
          maybeSum _ Nothing = Nothing
          maybeSum (Just a) (Just b) = Just (a+b)


isComplete :: Graph a -> Bool
isComplete g = all (\(a,b) -> a == b || areAdjacent g a b) 
                   [(i,j)| i <- [1..size g], j <- [1..size g]]

chromaticNumber :: Graph a -> Int
chromaticNumber graph = 
    chromaticNumber' graph $ 
        filter (\(a,b) -> a /= b && (not $ areAdjacent graph a b)) 
               [(i,j)| i <- [1..size graph], j <- [1..size graph]]
    where chromaticNumber' g [] = size g
          chromaticNumber' g ((a,b):es) = 
            min (chromaticNumber' (addBidirectEdge g a b undefined) es)
                (chromaticNumber $ mergeVertexes g a b) -- потому что список несмежных изменился
          mergeVertexes g i j = 
            let adjacents = getAdjacents g i ++ getAdjacents g j
                g' = flip delVertex j $ delVertex g i
                i' = size g'
            in L.foldl (\g'' k -> updateBidirectEdge g'' i' k undefined) 
                     g' adjacents

findPathes :: Graph a -> Int -> Int -> [[(Int,Int)]]
findPathes g s t 
        | s < 1 || t < 1 || 
          s > size g || t > size g || 
          (null $ getAdjacents g s) = []
        | otherwise = findPathes' $ getAdjacents g s
    where findPathes' = 
              filter (not . null) . 
              concatMap (\j -> explore [(s,j)] (delEdges g s j) ) 
          explore path@((i,j):es) g' = 
              if j == t 
               then [path] 
               else concatMap (\k -> explore [(j,k)] (delEdges g' j k)) $ 
                    getAdjacents g' j

maxFlow :: Graph Int -> Int -> Int -> Int
maxFlow g s t  
        | s < 1 || t < 1 || s > size g || t > size g = -1
        | otherwise = L.foldl (\s (f,_) -> if f > 0 then s + f else s) 0 $ 
                      mapMaybe (getWeight maxFlow' s) $ 
                      getAdjacents maxFlow' s
    where g' = L.foldl (\g' i -> 
                        L.foldl (\g'' j -> 
                                 if areAdjacent g'' j i 
                                  then g'' 
                                  else addEdge g'' j i (0,0)) 
                        g' $ getAdjacents g' i)
                     (mapEdges g (\(_,_,v) -> Just (0,v))) 
                     [i|i <- [1..size g]] -- скатать определение
          pathes = findPathes g' s t -- :: [(Int,Int)]
          maxFlow' = L.foldl activatePath g' pathes
          activatePath g'' path = 
              let localFlow = minimum $ map (freeFlow g'') path
              in L.foldl (addFlow localFlow) g'' path
          freeFlow g'' (i,j) = 
              case getWeight g'' i j of
               Nothing -> error "Some logical error in algo"
               Just (f,c) -> if c < 0 -- что делать при двунаправленных рёбрах? 
                              then freeFlow g'' (j,i) 
                              else c - f
          addFlow flow g'' (i,j) = 
              case getWeight g'' i j of
               Nothing -> error "Some logical error in algo"
               Just (f,c) -> if c < 0 
                              then updateEdge 
                                    (updateEdge g'' i j (f-flow,c)) 
                                    j i (-f+flow,-c)
                              else updateEdge 
                                    (updateEdge g'' i j (f+flow,c)) 
                                    j i (-f-flow,-c)


